package gupao.concurrency.threadpool;

import gupao.concurrency.enums.TrafficChannelEnum;
import gupao.concurrency.model.ForecastConfig;
import gupao.concurrency.model.ForcastContext;
import gupao.concurrency.model.TimeWindow;
import gupao.concurrency.service.ForcastConfigService;
import gupao.concurrency.service.TrafficForecastService;
import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

@Slf4j
public class ForecastScheduledExecutor extends ScheduledThreadPoolExecutor {

    private final ForcastConfigService configService;
    private final TrafficForecastService forecasterService;

    private final ConcurrentHashMap<ScheduledFuture<?>, ForecastRunnableTask> taskMap = new ConcurrentHashMap<>();
    private final ConcurrentHashMap<TrafficChannelEnum,ScheduledFuture<?>> channelMap = new ConcurrentHashMap<>();

    @Override
    protected void beforeExecute(Thread t, Runnable r) {
        super.beforeExecute(t, r);
        ScheduledFuture<?> sf = (ScheduledFuture<?>)r;
        ForecastRunnableTask task = taskMap.get(sf);
        if(task == null){
            log.warn("dirty task, clean");
            sf.cancel(true);
            return;
        }
        long periodInMillis = (long) task.getForecastConfig().get().getTimeWindowSeconds() * 1000;
        TimeWindow timeWindow = configService.createNewTimeWindow(task.getChannelEnum(), periodInMillis);
        ForcastContext.init(task.getForecastConfig().get(), timeWindow);
        task.beforeExec();
    }

    @Override
    protected void afterExecute(Runnable r, Throwable t) {
        super.afterExecute(r, t);
        ScheduledFuture<?> sf = (ScheduledFuture<?>)r;
        ForecastRunnableTask task = taskMap.get(sf);
        try{
            if(task == null){
                log.warn("dirty task");
                sf.cancel(true);
                return;
            }
            task.afterExec();
            boolean needRescheduleNewTask = false;
            if(task.getStatus().get() == ForecastRunnableTask.STATUS_WAIT_CANCEL){
                cancel(task.getChannelEnum(), true);
                if(task.getNextForecastConfig().get() != null){
                    needRescheduleNewTask = true;
                }
            }
            if(task.getStatus().get() == ForecastRunnableTask.STATUS_CANCELED){
                taskMap.remove(sf);
                channelMap.remove(task.getChannelEnum());
            }
            if(needRescheduleNewTask){
                scheduleForecastTask(task.getForecastConfig().get(), task.getChannelEnum());
            }
        }catch (Exception e){
            log.error("error in afterExec.",e);
        }finally {
            ForcastContext.clean();
            if(task != null){
                task.setStatus(ForecastRunnableTask.STATUS_AFTER_EXEC);
            }
        }

    }

    public ForecastScheduledExecutor(int corePoolSize,
                                     ForcastConfigService configService,
                                     TrafficForecastService forcasterService) {
        super(corePoolSize, new ThreadFactory() {
            AtomicInteger idx = new AtomicInteger(1);
            @Override
            public Thread newThread(Runnable r) {
                Thread t = new Thread(r);
                t.setDaemon(true);
                t.setName("traffic-forcast-thread-" + idx.getAndIncrement());
                return t;
            }
        });
        this.forecasterService = forcasterService;
        this.configService = configService;
    }

    public ForecastRunnableTask scheduleForecastTask(ForecastConfig config, TrafficChannelEnum channelEnum){
        ForecastRunnableTask task = new ForecastRunnableTask(channelEnum, config, forecasterService);
        ScheduledFuture<?> sf = scheduleAtFixedRate(task, config.getTimeWindowSeconds() * 1000);
        return sf == null ? null : task;
    }

    private ScheduledFuture<?> scheduleAtFixedRate(ForecastRunnableTask task, long periodInMillis){
        ScheduledFuture<?> sf = super.scheduleAtFixedRate(task, 1, periodInMillis, TimeUnit.MILLISECONDS);
        taskMap.put(sf, task);
        channelMap.put(task.getChannelEnum(), sf);
        return sf;
    }

    public void cancel(TrafficChannelEnum channel, boolean mayInterruptIfRunning){
        ScheduledFuture<?> sf = channelMap.get(channel);
        if(sf == null){
            log.warn("No task found for channel: {}", channel.name());
        } else {
            ForecastRunnableTask task = taskMap.get(channel);
            if(task != null){
                task.setStatus(ForecastRunnableTask.STATUS_CANCELED);
            }
            sf.cancel(mayInterruptIfRunning);
        }
    }

    public ForecastRunnableTask getRunningTask(TrafficChannelEnum channelEnum){
        ScheduledFuture<?> sf = channelMap.get(channelEnum);
        return sf == null ? null : taskMap.get(sf);
    }


    public void markCancel(TrafficChannelEnum channel){
        markCancelAndRescheduleLater(channel,null);
    }

    public void rescheduleAfterCancel(TrafficChannelEnum channel, ForecastConfig config){
        markCancelAndRescheduleLater(channel, config);
    }

    private void markCancelAndRescheduleLater(TrafficChannelEnum channel, ForecastConfig nextConfig){
        ScheduledFuture<?> sf = channelMap.get(channel);
        if(sf == null){
            return;
        }
        ForecastRunnableTask task = taskMap.get(channel);
        if(task != null){
            task.setStatus(ForecastRunnableTask.STATUS_WAIT_CANCEL);
            task.getNextForecastConfig().set(nextConfig);
        } else {
            sf.cancel(true);
            channelMap.remove(channel);
            scheduleForecastTask(nextConfig, channel);
        }
    }
}
